"""
https://leetcode.cn/problems/kth-largest-element-in-an-array/description/

215. 数组中的第K个最大元素
已解答
中等
相关标签
相关企业
给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。

 

示例 1:

输入: [3,2,1,5,6,4], k = 2
输出: 5
示例 2:

输入: [3,2,3,1,2,4,5,5,6], k = 4
输出: 4
 

提示：

1 <= k <= nums.length <= 105
-104 <= nums[i] <= 104

"""

from typing import List


class Solution:

    def swap(self,nums,i1,i2):
        nums[i1],nums[i2]=nums[i2],nums[i1]

    def cons_stack(self,sta:List[int])->List[int]:
        """
        构建堆栈
        父级元素（i）小于子集元素(i*2+1,i*2+2)
        子元素（n）的父元素((n-1)//2)
        """
        res=[]
        for i in sta:
            res.append(i)
            cidx=len(res)-1
            while cidx>0:
                pidx=(cidx-1)//2
                pval=res[pidx]
                if pval>i:
                    self.swap(res,pidx,cidx)
                    cidx=pidx
                else:
                    break
        return res

    def insert_new(self,sta:List[int],num:int)->None:
        """
        小顶堆插入新的元素
        """
        if sta[0]>num:
            return
        
        sta[0]=num
        pidx=0
        while True:
            cidx1=pidx*2+1
            cidx2=pidx*2+2
            if cidx2<len(sta):
                if sta[cidx1]<sta[cidx2]:
                    if sta[cidx1]<sta[pidx]:
                        self.swap(sta,cidx1,pidx)
                        pidx=cidx1
                    else:
                        break
                else:
                    if sta[cidx2]<sta[pidx]:
                        self.swap(sta,cidx2,pidx)
                        pidx=cidx2
                    else:
                        break
            elif cidx1<len(sta):
                if sta[cidx1]<sta[pidx]:
                    self.swap(sta,cidx1,pidx)
                    pidx=cidx1
                else:
                    break
            else:
                break


    def findKthLargest(self, nums: List[int], k: int) -> int:
        """
        典型的top-n问题，使用定数小顶堆解决
        """
        # 将前K个元素堆化
        sta=self.cons_stack(nums[:k])


        # 遍历之后的元素，并与堆顶元素对比，构成新的堆
        for i in range(k,len(nums)):
            self.insert_new(sta,nums[i])


        return sta[0]

if __name__=='__main__':
    sol=Solution()
    nums=[3,2,3,1,2,4,5,5,6]
    k = 4
    res=sol.findKthLargest(nums,k)
    print(res)